import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gym
from env_masterworkser import LHPCSEnv
import matplotlib.pyplot as plt

#=============== Hyper Parameters ====================================
BATCH_SIZE = 32              
LR = 0.001                     #learning rate
EPSILON = 0.9                 # e-greedy
GAMMA = 0.9                   # reward discount
TARGET_REPLACE_ITER = 100     # target update frequency
MEMORY_CAPACITY = 2000        
# env = gym.make('CartPole-v0') 
# env = env.unwrapped           # without time limit
env = LHPCSEnv(20)
N_ACTIONS = env.action_space.n  
N_STATES = env.observation_space.shape[0] * env.observation_space.shape[1] 
ENV_A_SHAPE = 0 if isinstance(env.action_space.sample(), int) else env.action_space.sample().shape     # to confirm the shape
TRAIN_CNT = 1500
#====================== Net Class ==========================================
class Net(nn.Module):
    def __init__(self, ):
        super(Net, self).__init__()
        # 这里以一个动作为作为观测值进行输入，然后把他们输出给50个神经元
        # N_STATES 与 图像的特征值个数有关
        self.fc1 = nn.Linear(N_STATES, 12)
        # N_ACTIONS 与 能做的动作个数有关
        self.fc1.weight.data.normal_(0, 0.1)   # 初始化权重，用二值分布来随机生成参数的值
        # 经过50个神经元运算过后的数据， 把每个动作的价值作为输出。
        self.out = nn.Linear(12, N_ACTIONS)    # 做出每个动作后，每个动作的价值作为输出。
        self.out.weight.data.normal_(0, 0.1)   # 初始化权重，用二值分布来随机生成参数的值
        # 输入-当前状态 action --Net网络--输出--》 所有动作价值
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        actions_value = self.out(x)
        return actions_value


#==================== DQN ================================================
net = Net()
class DQN(object):
    def __init__(self):
        #Q_evaluation, Q_target
        self.eval_net, self.target_net = Net(), Net()
        self.learn_step_counter = 0                                     # for target updating
        self.memory_counter = 0                                         # for storing memory
        #memory: s, a, r, s_
        self.memory = np.zeros((MEMORY_CAPACITY, N_STATES * 2 + 2))     # initialize memory
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=LR)
        self.loss_func = nn.MSELoss()

    def choose_action(self, x):
        x = torch.unsqueeze(torch.FloatTensor(x), 0)
        x = x.flatten()
        if np.random.uniform() < EPSILON:  
            # e-greedy
            actions_value = self.eval_net.forward(x) 
            actions_value = actions_value.data
            _, action = torch.max(actions_value, 0)
        else:  
            action = np.random.randint(0, N_ACTIONS)
            action = action if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE)
        return int(action)

    def store_transition(self, s, a, r, s_):
        # use numpy to store the data
        transition = np.hstack((s, [a, r], s_))
        index = self.memory_counter % MEMORY_CAPACITY
        self.memory[index, :] = transition
        self.memory_counter += 1

    def learn(self):
        if self.learn_step_counter % TARGET_REPLACE_ITER == 0:
            #eval -> target Net
            self.target_net.load_state_dict(self.eval_net.state_dict()) 
        self.learn_step_counter += 1
        sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE)  
        b_memory = self.memory[sample_index, :]  #batch_data
        #data -> s, a, r, s_
        b_s = torch.FloatTensor(b_memory[:, :N_STATES])  
        b_a = torch.LongTensor(b_memory[:, N_STATES:N_STATES+1].astype(int)) 
        b_r = torch.FloatTensor(b_memory[:, N_STATES+1:N_STATES+2])  
        b_s_ = torch.FloatTensor(b_memory[:, -N_STATES:]) 

        # q_eval w.r.t the action in experience
        #Q_eval(s, a)
        q_eval = self.eval_net(b_s).gather(1, b_a)  # shape (batch, 1)
        q_next = self.target_net(b_s_).detach()     # detach from graph, don't backpropagate
        # Q(s, a) = r + gamma * Q(s_, a_)
        q_target = b_r + GAMMA * q_next.max(1)[0].view(BATCH_SIZE, 1)   # shape (batch, 1)
        #optimize the loss between q_eval and q_target
        loss = self.loss_func(q_eval, q_target)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

#====================== Main ===========================
dqn = DQN()  
print('\nCollecting experience...')
x = []
y = []
# start training ...
for i_episode in range(TRAIN_CNT):
    s = env.reset() 
    ep_r = 0  #training accounter
    # loop will break only by the env
    while True:
        # env.render()  #visiualize
        a = dqn.choose_action(s)
        # next state: s
        # reward: r (float)
        # if done is true, env is over.
        # debug_information: info (no use)
        s_, r, done, info = env.step(a)
        # evaluate and revise
        # store data
        s = np.hstack(s)
        s_ = np.hstack(s_)
        dqn.store_transition(s, a, r, s_)
        ep_r += r
        if dqn.memory_counter > MEMORY_CAPACITY:
            # 它就会开对去学习。
            # eavl 每学一次就会更新一次  # 它的更新思路是从我历史记忆中随机抽取数据。 #学习一次，就在数据库中随机挑选BATCH_SIZE（32条） 进行打包

            # 而target不一样，它是在我们学习过程中到一定频率（TARGET_REPLACE_ITER，来决定）。它的思路是：target网会去复制eval网的参数
            dqn.learn()
            # 在满足 大于数据库容量的条件下，我再看env.step(a) 返回的done，env是否认为实验结束了
            if done:
                #print episode information
                if i_episode % 100 == 0:
                    print('Ep: ', i_episode,
                            '| Ep_r: ', ep_r)
                    print(env.actions)

                x.append(i_episode)
                y.append(ep_r)
        if done:
            break
        # next step, next epoch
        s = s_ 
#=========== draw =================
plt.plot(x, y)
plt.xlabel('counter')
plt.ylabel('reward')
plt.savefig('./result.png')